#从视频流获取截取帧
#从帧获取人脸图片并存储到对应文件夹
import cv2
import numpy as np
import dlib#这个库需要使用vs c++环境编译后才能安装成功
import os
import sys
import random
# 存储位置
output_dir = 'face_images'
origin_dir = 'cutted_frames'

size = 128#设置为128*128大小

if not os.path.exists(output_dir):
    os.makedirs(output_dir)
if not os.path.exists(origin_dir):
    os.makedirs(origin_dir)

#使用dlib自带的frontal_face_detector作为我们的特征提取器
detector = dlib.get_frontal_face_detector()

# 使用训练好的68个特征点模型
predictor_path = "shape_predictor_68_face_landmarks.dat"
predictor = dlib.shape_predictor(predictor_path)

# 打开摄像头 参数为输入流，可以为摄像头或视频文件
#此处选用了视频文件
#camera = cv2.VideoCapture(0) #使用摄像头
cap = cv2.VideoCapture('test.mp4')

#获得当前视频的帧率
#这个数值并不是整数，而是浮点数，因此可以采用int进行向下取整
#如果需要四舍五入用round，向上取整用ceil
FPS = cap.get(5)

timeRate = 5#截取视频帧的时间间隔，10即为每十秒截取1帧
counter = 0#设置计时器
index = 0

while(True):
    ret,frame = cap.read()
    if ret:
        frameRate = int(FPS)*timeRate
        if (counter % frameRate == 0):
            print("截取第: " + str(counter) + "帧")
            cv2.imwrite(origin_dir+'/'+str(counter)+'.jpg',frame)
        counter = counter + 1
    else:
        print("所有帧都已保存完成")
        # 释放摄像头 release camera
        cap.release()
        break

for filenames in os.listdir(origin_dir):
    print(filenames)
    img = cv2.imread(origin_dir + '/' +filenames)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    #使用detector检测器进行人脸检测，检测图像中的人脸
    faces = detector(gray, 1)
    print("人脸数: ", len(faces))
    for i, face in enumerate(faces):
        x1 = face.top() if face.top() > 0 else 0
        y1 = face.bottom() if face.bottom() > 0 else 0
        x2 = face.left() if face.left() > 0 else 0
        y2 = face.right() if face.right() > 0 else 0

        face = img[x1:y1,x2:y2]
        #重新调整人脸照片的大小
        face = cv2.resize(face, (size,size))
        cv2.imshow('image', face)
        cv2.imwrite(output_dir+'/'+str(index)+'.jpg', face)
        index += 1